import sys
sys.path.append('../../')
from BasicLibraries import *
import allocation_ga as af
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning) 
warnings.filterwarnings("ignore", category=FutureWarning) 


AM = af.AllocationGAD()

def create_landscape_files(FW, dt, I,  env_sdgs, tmp_outputs, parent_selection_option, crossover_option,\
                           mutation_types_option, normal_fit=True, allocation_type = 'GA', final_year=2021,energy_control=False):
    #==
    dt = dt.groupby('mrg').last()
    dt = dt[dt.next_year_performance > -3]
    dt = dt.sort_values(by = ['gvkey', 'rfyear']).reset_index()
    s = dt.copy()
    s = s[(s.rfyear >= 2015)]
    idx = s[['gvkey', 'rfyear']].groupby('gvkey').count()
    idx = idx[idx >= 3].dropna()
    dt = dt[dt.gvkey.isin(idx.index)]

    #==
    gvkey_list = list(np.unique(dt[dt.gvkey.isin(idx.index)].gvkey))
    print('Number of firms to analyse:', len(gvkey_list))
    sdg_universe = 'Environment'
    #== Create the landscape
    max_year = final_year + 1 # The +1 is simply for the boundary of the for loop
    min_year = max_year - 4

    for Y in range(min_year,max_year):
        #== Only keep firms with an observation in year Y
        reduced_dt = dt[dt.rfyear <= Y].copy()
        idx = reduced_dt[['gvkey', 'rfyear']].groupby('gvkey').last()
        reduced_dt = reduced_dt[reduced_dt.gvkey.isin(idx[idx.rfyear == Y].dropna().index)]
        #== Run
        resultsDICTIONARY, relevant_features, irrelevant_features = AM.main_run(reduced_dt, \
                                        parent_selection_option,crossover_option,mutation_types_option, \
                                        sdgs_number= env_sdgs,  
                                        dimension_ = I, 
                                        gvkey_list = gvkey_list, 
                                        model_type = 'RandomForest', \
                                        quantile_bound= 0.99, 
                                        performance_percentile = 0.95, 
                                        normal_fit = normal_fit, energy_control=energy_control,FW=FW)

        company_score, difference_mat, optima, observed = [], [], [], []
        difference_mat_dict, optima_dict, observed_dict = dict(), dict(), dict()
        for i in range(len(resultsDICTIONARY)):
            try:
                g, obs, opt, difference_matrix, scores_stats, g_score_vec, g_score, distance_from_closest_optimum  = \
                    AM.calculate_firm_statistics(resultsDICTIONARY, i,  reduced_dt)
                company_score.append([g, g_score, distance_from_closest_optimum ])
                difference_mat.append(difference_matrix)
                optima.append(opt)
                observed.append(obs)
                difference_mat_dict[g] = difference_matrix
                optima_dict[g] = opt
                observed_dict[g] = obs
            except Exception:
                print(i, 'Error')
                continue
        company_score = pd.DataFrame(company_score, columns = ['gvkey', 'score', 'distance_from_opt'])
        names = reduced_dt[['conml', 'gvkey', 'loc']].groupby('gvkey').last().reset_index()
        reduced_dt['Tangibility'] = reduced_dt['ppent_usd']/reduced_dt['at_usd']
        reduced_dt['buyback_rate'] = (reduced_dt['iss_eq_usd_ref']/(reduced_dt.sale_usd*1000000))
        characteristics = reduced_dt[['gvkey', 'Total_invested_capital_BS1', 'firm_size', 'iss_eq_usd_ref', 'dps_usd_ref', 'buyback_rate', 'Tangibility']].groupby('gvkey').mean().reset_index()
        company_score = company_score.merge(names).merge(characteristics)
        company_score['rfyear'] = [Y]*len(company_score)


        pickle.dump([resultsDICTIONARY,relevant_features, irrelevant_features,\
                     company_score, observed, \
                     optima, observed_dict, optima_dict, gvkey_list, reduced_dt], open(tmp_outputs+str(Y)+'.pckl', 'wb'))

